"""会议安排任务的Agent Loop实现"""
import asyncio
import json
import logging
import os
import copy
from typing import Any
from uuid import uuid4

from jinja2 import Template
from tokenizers import Tokenizer
from verl.experimental.agent_loop.agent_loop import AgentLoopBase, AgentLoopOutput, register
from verl.utils.profiler import simple_timer
from verl.utils.rollout_trace import rollout_trace_op

from agentic_system.environments.prompts.meeting_prompt import SYSTEM_TEMPLATE, MEETING_TEMPLATE
from agentic_system.environments.env_package.meeting.meeting.meeting_agent_env import MeetingScheduleEnv
import agentic_system.environments.env_package.meeting.projection as proj
from agentic_system.environments.env_package.meeting.meeting_reward import MeetingReward

logger = logging.getLogger(__file__)
logger.setLevel(os.getenv("VERL_LOGGING_LEVEL", "WARN"))


@register("meeting_agent")
class MeetingAgentLoop(AgentLoopBase):
    """会议安排任务的Agent Loop"""
    
    @classmethod
    def init_class(cls, config, tokenizer: Tokenizer, **kwargs):
        """类级别的初始化"""
        if cls._class_initialized:
            return
        cls._class_initialized = True
        print("正在初始化 MeetingAgentLoop 类...")
        
        # 初始化tokenizer和配置
        cls.tokenizer = tokenizer
        cls.max_user_turns = config.actor_rollout_ref.rollout.multi_turn.max_user_turns
        cls.max_assistant_turns = config.actor_rollout_ref.rollout.multi_turn.max_assistant_turns
        cls.prompt_length = config.actor_rollout_ref.rollout.prompt_length
        cls.response_length = config.actor_rollout_ref.rollout.response_length
        
        # 初始化环境和奖励
        cls.env = MeetingScheduleEnv(max_steps=cls.max_assistant_turns)
        cls.reward_calculator = MeetingReward()
        
        # 初始化日志文件
        if config.trainer.validation_data_dir:
            log_file_name = os.path.join(config.trainer.validation_data_dir, f"{uuid4().hex}.log")
            if not os.path.exists(config.trainer.validation_data_dir):
                os.makedirs(config.trainer.validation_data_dir, exist_ok=True)
            cls.log_file = open(log_file_name, "w", encoding='utf-8')
            print(f"Agent loop 日志文件: {log_file_name}")
        else:
            cls.log_file = None
        
        print(f"MeetingAgentLoop 初始化完成")
        print(f"- 最大assistant轮数: {cls.max_assistant_turns}")
        print(f"- 最大user轮数: {cls.max_user_turns}")
        print(f"- Prompt长度: {cls.prompt_length}")
        print(f"- Response长度: {cls.response_length}")
    
    @rollout_trace_op
    async def run(self, sampling_params: dict[str, Any], **kwargs):
        """
        运行一次完整的会议安排任务
        
        Args:
            sampling_params: 采样参数
            **kwargs: 任务数据，包括：
                - meeting_topic: 会议主题
                - required_attendees: 必选人员列表
                - optional_attendees: 可选人员列表
                - availability: 人员日程可用性
                - historical_pattern: 历史会议模式
        """
        # 提取任务数据
        task_data = {
            'meeting_topic': kwargs.get('meeting_topic', ''),
            'required_attendees': kwargs.get('required_attendees', []),
            'optional_attendees': kwargs.get('optional_attendees', []),
            'availability': kwargs.get('availability', {}),
            'historical_pattern': kwargs.get('historical_pattern', {})
        }
        
        # 重置环境
        observation, info = await self.loop.run_in_executor(
            None,
            lambda: self.env.reset(task_data)
        )
        
        # 初始化状态
        user_turns = 0
        assistant_turns = 0
        total_reward = 0
        mt_res = []  # 多轮结果列表
        history_text = ""
        task_id = uuid4().hex
        
        if self.log_file:
            self.log_file.write(f"\n{'='*80}\n")
            self.log_file.write(f"任务ID: {task_id}\n")
            self.log_file.write(f"会议主题: {task_data['meeting_topic']}\n")
            self.log_file.write(f"必选人员: {', '.join(task_data['required_attendees'])}\n")
            self.log_file.write(f"可选人员: {', '.join(task_data['optional_attendees'])}\n")
            self.log_file.write(f"{'='*80}\n\n")
        
        # 主循环
        while True:
            response_mask = []
            messages = []
            metrics = {}
            
            # 构建消息
            messages.append({"role": "system", "content": SYSTEM_TEMPLATE})
            
            # 构建用户提示词
            user_prompt = Template(MEETING_TEMPLATE).render(
                meeting_topic=task_data['meeting_topic'],
                required_attendees=', '.join(task_data['required_attendees']),
                optional_attendees=', '.join(task_data['optional_attendees']),
                observation=observation,
                history=history_text
            )
            messages.append({"role": "user", "content": user_prompt})
            
            # Tokenize提示词
            request_id = uuid4().hex
            prompt_ids = await self.loop.run_in_executor(
                None,
                lambda: self.tokenizer.apply_chat_template(
                    messages, add_generation_prompt=True, tokenize=True
                ),
            )
            
            # 截断过长的prompt
            if len(prompt_ids) > self.prompt_length:
                prompt_ids = prompt_ids[-self.prompt_length:]
            
            # 生成响应
            with simple_timer("generate_sequences", metrics):
                output = await self.server_manager.generate(
                    request_id=request_id,
                    prompt_ids=prompt_ids,
                    sampling_params=sampling_params,
                    image_data=None
                )
            
            response_ids = output.token_ids
            
            # 解码响应
            response = await self.loop.run_in_executor(
                None,
                lambda: self.tokenizer.decode(
                    response_ids,
                    skip_special_tokens=True,
                    clean_up_tokenization_spaces=True
                ),
            )
            
            response_mask += [1] * len(response_ids)
            assistant_turns += 1
            
            if self.log_file:
                self.log_file.write(f"\n--- 第 {assistant_turns} 轮 ---\n")
                self.log_file.write(f"Prompt长度: {len(prompt_ids)}\n")
                self.log_file.write(f"Response: {response}\n")
            
            # 检查终止条件
            # 1. 达到最大response长度
            if len(response_mask) >= self.response_length:
                if self.log_file:
                    self.log_file.write(f"终止原因：达到最大response长度 {len(response_mask)}\n")
                break
            
            # 2. 达到最大assistant轮数
            if self.max_assistant_turns and assistant_turns >= self.max_assistant_turns:
                if self.log_file:
                    self.log_file.write(f"终止原因：达到最大assistant轮数 {assistant_turns}\n")
                break
            
            # 3. 达到最大user轮数
            if self.max_user_turns and user_turns >= self.max_user_turns:
                if self.log_file:
                    self.log_file.write(f"终止原因：达到最大user轮数 {user_turns}\n")
                break
            
            # 解析动作
            action_lst, invoke_lst, valid_lst = proj.meeting_projection([response])
            action = action_lst[0]
            invoke = invoke_lst[0]
            valid = valid_lst[0]
            
            if self.log_file:
                self.log_file.write(f"解析结果 - Valid: {valid}, Action: {action}, Invoke: {invoke}\n")
            
            # 检查动作有效性
            if not valid:
                if self.log_file:
                    self.log_file.write(f"终止原因：动作格式无效\n")
                break
            
            # 执行动作
            with simple_timer("environment_step", metrics):
                obs_result, step_reward, is_done, step_info = await self.loop.run_in_executor(
                    None,
                    lambda: self.env.step(invoke, valid)
                )
            
            total_reward += step_reward
            observation = obs_result
            
            if self.log_file:
                self.log_file.write(f"环境反馈: {obs_result}\n")
                self.log_file.write(f"步骤奖励: {step_reward}, 累计奖励: {total_reward}\n")
            
            # 更新历史
            history_text += f"\n轮次 {assistant_turns}:\n"
            history_text += f"动作: {action.get('description', '')}\n"
            history_text += f"工具调用: {invoke.get('tool_name', '')}\n"
            history_text += f"结果: {obs_result[:200]}...\n"  # 截断过长的结果
            
            # 保存当前轮次结果
            output_obj = AgentLoopOutput(
                prompt_ids=prompt_ids if len(prompt_ids) <= self.prompt_length else prompt_ids[-self.prompt_length:],
                response_ids=response_ids if len(response_ids) <= self.response_length else response_ids[:self.response_length],
                response_mask=response_mask if len(response_mask) <= self.response_length else response_mask[:self.response_length],
                num_turns=user_turns + assistant_turns + 1,
                metrics=metrics,
                reward_score=total_reward,
                extra_fields=dict()
            )
            mt_res.append(output_obj)
            
            # 检查任务是否完成
            if is_done:
                if self.log_file:
                    self.log_file.write(f"任务完成！最终奖励: {total_reward}\n")
                break
            
            user_turns += 1
        
        # 确保至少有一个输出
        if not mt_res:
            output_obj = AgentLoopOutput(
                prompt_ids=prompt_ids if len(prompt_ids) <= self.prompt_length else prompt_ids[-self.prompt_length:],
                response_ids=response_ids if len(response_ids) <= self.response_length else response_ids[:self.response_length],
                response_mask=response_mask if len(response_mask) <= self.response_length else response_mask[:self.response_length],
                num_turns=user_turns + assistant_turns + 1,
                metrics=metrics,
                reward_score=total_reward,
                extra_fields=dict()
            )
            mt_res.append(output_obj)
        
        # 统一设置最终奖励和轮数
        final_reward = mt_res[-1].reward_score
        final_turns = mt_res[-1].num_turns
        
        for o in mt_res:
            o.reward_score = final_reward
            o.num_turns = final_turns
            o.extra_fields["count"] = len(mt_res)
        
        if self.log_file:
            self.log_file.write(f"\n总轮数: {len(mt_res)}, 最终奖励: {final_reward}\n")
            self.log_file.flush()
        
        return mt_res

